import os
# 图像数据的加载、处理和转换
import torchvision.transforms as transforms #调整尺寸、裁剪、归一化、旋转、翻转
import torchvision.datasets as datasets # 加载数据集
from torch.utils.data import DataLoader # 加载和批量处理大型数据集
from torchvision.datasets import ImageNet  # 大规模的图像分类数据集
from .utils import RandomSizedRectCrop, RectScale, RandomErasing, Random2DTranslation,WikiArtPolicy

class Data:
    def __init__(self, opt):
        pin_memory = False  # 是否加载到cuda的固定内存中

        scale_size = 256

        traindir = './database/WikiArt/train'
        valdir = './database/WikiArt/val'
        # 图像的像素标准化预处理
        normalize = transforms.Normalize((0.5071, 0.4865, 0.4409), (0.2673, 0.2564, 0.2762))

        # 加载训练集
        trainset = datasets.ImageFolder(
            traindir,
            transforms.Compose([
                transforms.Resize(scale_size),
                RandomSizedRectCrop(224, 224),
                # transforms.RandomResizedCrop(224, scale=(0.08, 1.0)),
                transforms.RandomHorizontalFlip(),
                transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 颜色抖动
                transforms.RandomRotation(15),  # 添加随机旋转
                RandomErasing(),# 随机擦除
                Random2DTranslation(224, 224, p=0.5),  # 随机二维平移
                # transforms.ToPILImage(),
                # WikiArtPolicy(),
                transforms.ToTensor(),
                normalize,
                # RectScale(224, 224),
            ]))

        # 创建训练集的数据加载器
        self.trainLoader = DataLoader(
            trainset,  # 要加载的数据集
            batch_size=opt.train_batch_size,  # 每个批次加载的数据量 128
            shuffle=True,  # 是否打乱
            num_workers=4,  # 使用4个子进程
            pin_memory=pin_memory, # 是否加载到cuda的固定内存中 从而加速数据传输到 GPU。
            drop_last=True,)  # 丢弃数据集中的最后一个不满 batch_size 的批次

        # 加载验证集
        testset = datasets.ImageFolder(
            valdir,
            transforms.Compose([
                transforms.Resize(scale_size),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                normalize,
            ]))

        # 创建验证集的数据加载器
        self.testLoader = DataLoader(
            testset,
            batch_size=opt.test_batch_size,
            shuffle=False,
            num_workers=4,
            pin_memory=pin_memory,
            drop_last=False,)

